module Parts

import Base: size, setindex!, getindex, @propagate_inbounds

abstract type AbstractParts{T, N, A, I} <: AbstractArray{T, N} end

Base.size(r::AbstractParts) = size(r.locations)
@propagate_inbounds Base.setindex!(r::AbstractParts{T, N}, v, index::Vararg{Int, N}) where {T, N} =
    r.array[r.locations[index...]...] = v

export Arrays
"""
    Arrays

A collection of arrays at certain positions within a larger array.

```jldoctest
julia> using Parts, Base.Test

julia> a = @inferred Arrays([1, 3, 2, 4], [(1:2,), (3:4,)]);

julia> @inferred a[1]
2-element Array{Int64,1}:
 1
 3

julia> a[1] = [3, 5];

julia> @inferred map(sum, a)
2-element Array{Int64,1}:
 8
 6
```
"""
struct Arrays{T, N, A, I} <: AbstractParts{T, N, A, I}
    array::A
    locations::I
end

@propagate_inbounds getindex(r::Arrays{T, N}, index::Vararg{Int, N}) where {T, N} =
    r.array[r.locations[index...]...]

Arrays(array::AbstractArray, locations::AbstractArray) =
    Arrays{
        typeof(array[first(locations)...]),
        ndims(locations),
        typeof(array),
        typeof(locations)
    }(array, locations)

export Views

"""
    Views

A collection of views at certain positions within a larger array.

```jldoctest
julia> using Parts, Base.Test

julia> a = @inferred Views([1, 3, 2, 4], [(1:2,), (3:4,)]);

julia> @inferred a[1]
2-element SubArray{Int64,1,Array{Int64,1},Tuple{UnitRange{Int64}},true}:
 1
 3

julia> @inferred map(sum, a)
2-element Array{Int64,1}:
 4
 6
```
"""
struct Views{T, N, A, I} <: AbstractParts{T, N, A, I}
    array::A
    locations::I
end

Views(array::AbstractArray, locations::AbstractArray) =
    Views{
        typeof(@view array[first(locations)...]),
        ndims(locations),
        typeof(array),
        typeof(locations)
    }(array, locations)

@propagate_inbounds getindex(r::Views{T, N}, index::Vararg{Int, N}) where {T, N} =
    @view r.array[r.locations[index...]...]

export Swaps
"""
    Swaps

A collection of swaps at certain positions within a larger array. Swaps must all
be the same size, and one container will be reused to access them.

```jldoctest
julia> using Parts, Base.Test

julia> a = @inferred Swaps([1, 3, 2, 4], [(1:2,), (3:4,)]);

julia> @inferred a[1]
2-element Array{Int64,1}:
 1
 3

julia> @inferred map(sum, a)
2-element Array{Int64,1}:
 4
 6

julia> b = @inferred Swaps([1, 3, 2, 4], [(1:3,), (4:4,)]);

julia> @inferred b[2]
ERROR: Swaps must all be the same size
[...]
```
"""
struct Swaps{T, N, A, I} <: AbstractParts{T, N, A, I}
    array::A
    locations::I
    swap::T
end

function Swaps(array::AbstractArray, locations::AbstractArray)
    swap = similar(array, size(@view array[first(locations)...])...)
    Swaps{
        typeof(swap),
        ndims(locations),
        typeof(array),
        typeof(locations)
    }(array, locations, swap)
end

@propagate_inbounds function getindex(s::Swaps{T, N}, index::Vararg{Int, N}) where {T, N}
    swap = s.swap
    array = s.array
    new_index = s.locations[index...]
    @boundscheck if size(swap) != size(@view array[new_index...])
        error("Swaps must all be the same size")
    end
    Base._unsafe_getindex!(swap, array, new_index...)
end

end
